#include <iostream>
#include <vector>
#include <string>
#include <Eigen/Dense>

#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/TwistStamped.h>
#include <mavros_msgs/AttitudeTarget.h>
#include <geometry_msgs/Vector3.h>
#include <geometry_msgs/Quaternion.h>
#include <mavros_msgs/State.h>
#include <nav_msgs/Odometry.h>
#include <sensor_msgs/Imu.h>

#include "offboard/ControlCommand.h"
#include "offboard/DebugPID.h"

#include "controller/pid.hpp"
// #include "controller/ThrustEstimator.hpp"
#include "controller/colors.h"

constexpr int LOOPRATE = 50; // TODO control loop rate
constexpr double RAD2ANGLE = 180.0 / M_PI;
constexpr double GRAVITATIONAL_ACC = 9.794;      // TODO： The gravity of Shanghai
constexpr double YAWTHRE = M_1_PI;               // TODO: 最大yaw角变化量(角速度)
constexpr double MAXTHETA = M_PI * 20.0 / 180.0; // TODO: max title angle
double MAXCOS = cos(MAXTHETA);

/* -------------------------------------------------------------------------- */
/*                               Global Variable                              */
/* -------------------------------------------------------------------------- */
mavros_msgs::State _cur_state;
// geometry_msgs::PoseStamped _pos;
// geometry_msgs::TwistStamped _vel;
sensor_msgs::Imu _imu;
Eigen::Vector3d _acc;
nav_msgs::Odometry _odom;
offboard::ControlCommand _pva_msg;

ros::Publisher _pubPID;
ros::Publisher _pubAttThrust;

double _hover_thrust = 0.4;           // 飞机基础油门值。各架飞机不同
vector<double> _thrust_limit(2, 0.0); // 油门上下限
vector<double> _posKP(3, 0.0);
vector<double> _posKI(3, 0.0);
vector<double> _posKD(3, 0.0);
vector<double> _velKP(3, 0.0);
vector<double> _velKI(3, 0.0);
vector<double> _velKD(3, 0.0);

PID3 _pid_pos, _pid_vel;

// ThrustEstimator _thrust_estimator;
// ThrustEstimator _thrust_estimator2;

bool _isRecord = true;
bool _isPrint = true;
bool _haveTarget = false;

/* -------------------------------------------------------------------------- */
/*                            Function Declaration                            */
/* -------------------------------------------------------------------------- */
// 回调函数
void cb_state(const mavros_msgs::State::ConstPtr &msg)
{
  _cur_state = *msg;
}

// void cb_pos(const geometry_msgs::PoseStamped::ConstPtr &msg)
// {
//     _pos = *msg;
// }

// void cb_vel(const geometry_msgs::TwistStamped::ConstPtr &msg)
// {
//     _vel = *msg;
// }

void cb_imu(const sensor_msgs::Imu::ConstPtr &msg)
{
  _imu = *msg;
  Eigen::Vector3d acc_body = Eigen::Vector3d(_imu.linear_acceleration.x,
                                             _imu.linear_acceleration.y,
                                             _imu.linear_acceleration.z);
  Eigen::Quaterniond quat(_imu.orientation.w,
                          _imu.orientation.x,
                          _imu.orientation.y,
                          _imu.orientation.z);
  _acc = quat * acc_body;
}

void cb_odom(const nav_msgs::Odometry::ConstPtr &msg)
{
  _odom = *msg;
}

void cb_target_pva(const offboard::ControlCommand::ConstPtr &msg)
{
  _pva_msg = *msg;
  _haveTarget = true;
}

double getYawFromQuaternion(geometry_msgs::Quaternion q)
{
  Eigen::Quaterniond quat(q.w, q.x, q.y, q.z);
  Eigen::Matrix3d R = quat.toRotationMatrix();
  return atan2(R(1, 0), R(0, 0));
}

/**
 * 将四元数转化为欧拉角形式，输入的顺序为x-y-z-w,欧拉角的顺序为Yaw(z)-Pitch(y)-Roll(x)
 */
Eigen::Vector3d quaternion2euler(double x, double y, double z, double w)
{
  Eigen::Vector3d ret;
  ret(0) = atan2(2.0 * (w * x + y * z), 1.0 - 2.0 * (x * x + y * y));
  ret(1) = asin(2.0 * (w * y - z * x));
  ret(2) = atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z));
  return ret;
}

Eigen::Vector3d quaternion2Euler_ros(geometry_msgs::Quaternion q)
{
  return quaternion2euler(q.x, q.y, q.z, q.w);
}

/**
 * @brief 从rosparam中的参数赋给PID
 */
bool setPIDParam()
{
  if (!_pid_pos.setPID(_posKP, _posKI, _posKD)) return false;
  if (!_pid_vel.setPID(_velKP, _velKI, _velKD)) return false;

  _pid_pos.setLimit({0.4, 0.4, 1.5}, // 积分上限
                    {2, 2, 3},       // 输出上限
                    {0, 0, 0});      // 死区
  _pid_vel.setLimit({0.2, 0.2, 1.5},
                    {2, 2, 3},
                    {0, 0, 0});

  return true;
}

/**
 * @brief 最重要的计算函数
 *
 * @param t_cur 当前时间
 * @param pos_cur   当前位置
 * @param vel_cur   当前速度
 * @param pos_tar 目标位置
 * @param vel_tar 目标速度
 * @param acc_tar 目标加速度
 * @param yaw_tar 目标yaw角度
 * @param state     无人机此时的状态，是否解锁，是否进入OFFBOARD
 */
void px4AttitudeCtlPVA(double t_cur,
                       const Eigen::Vector3d &pos_cur,
                       const Eigen::Vector3d &vel_cur,
                       const Eigen::Vector3d &pos_tar,
                       const Eigen::Vector3d &vel_tar,
                       const Eigen::Vector3d &acc_tar,
                       double yaw_tar,
                       const mavros_msgs::State &state)
{
  if (state.mode == "OFFBOARD")
  {
    // _thrust_estimator.setBeginFlag(true);
    // _thrust_estimator2.setBeginFlag(true);

    /// 位于较好的初始位置，再启用积分环节
    /// 如果处于Offboard状态再用积分环节
    if ((pos_cur - pos_tar).norm() < 1)
    {
      ROS_WARN_ONCE("[Controller] Start intergrate in PID");
      _pid_pos.setIntergrateFlag(true);
      _pid_vel.setIntergrateFlag(true);
    }
    // else {
    //   ROS_WARN_ONCE("[Controller] Stop intergrate in PID");
    //   _pid_pos.setIntergrateFlag(false);
    //   _pid_vel.setIntergrateFlag(false);
    // }
  }

  Eigen::Vector3d posError, velError;
  posError = pos_tar - pos_cur;
  velError = vel_tar - vel_cur;

  // calculate PID
  _pid_pos.add_error(posError, t_cur);
  _pid_vel.add_error(velError, t_cur);

  Eigen::Vector3d pos_out = _pid_pos.calculate();
  Eigen::Vector3d vel_out = _pid_vel.calculate();

  Eigen::Vector3d accExcept = pos_out + vel_out + acc_tar;

  accExcept[2] = min(1.5, max(accExcept[2], -1.5)); // TODO:限制Z方向加速度

  Eigen::Vector3d accExcept_g = accExcept +
                                Eigen::Vector3d(0, 0, GRAVITATIONAL_ACC);

  // Limit tilt angle to MAXTHETA
  if (Eigen::Vector3d(0, 0, 1).dot(accExcept_g.normalized()) < MAXCOS)
  {
    ROS_WARN_STREAM("Control reaches the limitation: " << accExcept_g.transpose());
    double nf_ = accExcept.norm();
    double A_ = MAXCOS * MAXCOS * nf_ * nf_ - accExcept(2) * accExcept(2);
    double B_ = 2 * (MAXCOS * MAXCOS - 1) * accExcept(2) * GRAVITATIONAL_ACC;
    double C_ = (MAXCOS * MAXCOS - 1) * GRAVITATIONAL_ACC * GRAVITATIONAL_ACC;
    double s = (-B_ + sqrt(B_ * B_ - 4 * A_ * C_)) / (2 * A_);
    accExcept_g.noalias() = s * accExcept +
                            GRAVITATIONAL_ACC * Eigen::Vector3d(0, 0, 1);
  }

  // Add control for yaw.
  Eigen::Vector3d b1c, b2c, b3c;
  Eigen::Vector3d b2d(-sin(yaw_tar), cos(yaw_tar), 0);
  if (accExcept_g.norm() > 1e-6)
    b3c.noalias() = accExcept_g.normalized();
  else
    b3c.noalias() = Eigen::Vector3d(0, 0, 1);

  b1c.noalias() = b2d.cross(b3c).normalized();
  b2c.noalias() = b3c.cross(b1c).normalized();
  Eigen::Matrix3d R_;
  Eigen::Quaterniond qua_;
  R_ << b1c, b2c, b3c;
  qua_ = Eigen::Quaterniond(R_);

  // pub
  double msg_thrust;
  geometry_msgs::Quaternion msg_qua;
  mavros_msgs::AttitudeTarget msg_AT;

  msg_thrust = accExcept_g.norm() / GRAVITATIONAL_ACC * (_hover_thrust);
  msg_thrust = min(max(msg_thrust, _thrust_limit[0]), _thrust_limit[1]);

  // _thrust_estimator.push(msg_thrust);
  // double esti_thrust = _thrust_estimator.getThrust(
  //     _imu.linear_acceleration.z, GRAVITATIONAL_ACC);

  // _thrust_estimator2.push(accExcept_g(2) / GRAVITATIONAL_ACC * (_hover_thrust));
  // double esti_thrust2 = _thrust_estimator2.getThrust(
  //     _acc(2), GRAVITATIONAL_ACC);

  msg_qua.w = qua_.w();
  msg_qua.x = qua_.x();
  msg_qua.y = qua_.y();
  msg_qua.z = qua_.z();

  msg_AT.orientation = msg_qua;
  msg_AT.thrust = msg_thrust;
  msg_AT.header.stamp = ros::Time::now();
  msg_AT.type_mask = mavros_msgs::AttitudeTarget::IGNORE_ROLL_RATE |
                     mavros_msgs::AttitudeTarget::IGNORE_PITCH_RATE |
                     mavros_msgs::AttitudeTarget::IGNORE_YAW_RATE;
  _pubAttThrust.publish(msg_AT);

  if (_isPrint) {
    Eigen::Vector3d angleTarget = quaternion2Euler_ros(msg_qua);

    cout.precision(3);
    cout << GREEN_B
         << "yaw: " << angleTarget(2)
         << " pitch: " << angleTarget(1)
         << " roll: " << angleTarget(0)
         << " thrust: " << msg_thrust << RESET << endl;

    cout << GREEN
         << "pos_cur  :" << pos_cur.transpose() << endl
         << "pos_tar:" << pos_tar.transpose() << endl
         << "posError :" << posError.transpose() << endl
         << "OUT pos  :" << pos_out.transpose() << endl
         << "vel_cur  :" << vel_cur.transpose() << endl
         << "vel_tar:" << vel_tar.transpose() << endl
         << "velError :" << velError.transpose() << endl
         << "OUT vel  :" << vel_out.transpose() << endl
         << "acc_tar:" << acc_tar.transpose() << endl
         << "accExcept:" << accExcept.transpose()
         << RESET << endl;
  }

  if (_isRecord) {
    Eigen::Vector3d angleTarget = quaternion2Euler_ros(msg_qua);
    offboard::DebugPID msgPID;
    msgPID.header.stamp = msg_AT.header.stamp;
    msgPID.pos_error.x = posError[0];
    msgPID.pos_error.y = posError[1];
    msgPID.pos_error.z = posError[2];
    msgPID.vel_error.x = velError[0];
    msgPID.vel_error.y = velError[1];
    msgPID.vel_error.z = velError[2];
    msgPID.pos_pid_out.x = pos_out[0];
    msgPID.pos_pid_out.y = pos_out[1];
    msgPID.pos_pid_out.z = pos_out[2];
    msgPID.vel_pid_out.x = vel_out[0];
    msgPID.vel_pid_out.y = vel_out[1];
    msgPID.vel_pid_out.z = vel_out[2];
    msgPID.acc_except.x = accExcept[0];
    msgPID.acc_except.y = accExcept[1];
    msgPID.acc_except.z = accExcept[2];
    msgPID.euler_out.x = angleTarget[0] /* * RAD2ANGLE */;
    msgPID.euler_out.y = angleTarget[1] /* * RAD2ANGLE */;
    msgPID.euler_out.z = angleTarget[2] /* * RAD2ANGLE */;
    // msgPID.hover_thrust = esti_thrust;
    _pubPID.publish(msgPID);
  }
}

/* -------------------------------------------------------------------------- */
/*                                main function                               */
/* -------------------------------------------------------------------------- */
int main(int argc, char **argv)
{
  // node init
  ros::init(argc, argv, "control_node");
  ros::NodeHandle nh;
  ros::Rate rate(LOOPRATE);

  // load parameter file
  const string &node_name = ros::this_node::getName();
  nh.getParam(node_name + "/" + "BaseThrust", _hover_thrust);
  nh.getParam(node_name + "/" + "ThrustLimit", _thrust_limit);
  nh.getParam(node_name + "/" + "posKP", _posKP);
  nh.getParam(node_name + "/" + "posKI", _posKI);
  nh.getParam(node_name + "/" + "posKD", _posKD);
  nh.getParam(node_name + "/" + "velKP", _velKP);
  nh.getParam(node_name + "/" + "velKI", _velKI);
  nh.getParam(node_name + "/" + "velKD", _velKD);
  nh.getParam(node_name + "/" + "isRecord", _isRecord);
  nh.getParam(node_name + "/" + "isPrint", _isPrint);
  std::cout
      << GREEN_B
      << "BaseThrust: " << _hover_thrust << endl
      << "ThrustLimit:" << _thrust_limit[0] << "  ~  " << _thrust_limit[1] << endl
      << "posKP:      " << _posKP[0] << " " << _posKP[1] << " " << _posKP[2] << endl
      << "velKP:      " << _velKP[0] << " " << _velKP[1] << " " << _velKP[2] << endl
      << "velKI:      " << _velKI[0] << " " << _velKI[1] << " " << _velKI[2] << endl
      << "velKD:      " << _velKD[0] << " " << _velKD[1] << " " << _velKD[2]
      << RESET << endl;

  // Check thrust range
  if (_thrust_limit[1] < _hover_thrust + 0.05 ||
      _thrust_limit[0] > _hover_thrust - 0.05)
  {
    std::cout << RED_B
              << "The param `ThrustLimit` is unreasonable, EXIT !!!!!!!!"
              << RESET << endl;
    return 0;
  }

  // _thrust_estimator.init(GRAVITATIONAL_ACC / _hover_thrust);
  // _thrust_estimator2.init(GRAVITATIONAL_ACC / _hover_thrust);

  /// subscriber
  ros::Subscriber subState = nh.subscribe<mavros_msgs::State>(
      "mavros/state", 1, &cb_state);
  // ros::Subscriber subLocalPos = nh.subscribe<geometry_msgs::PoseStamped>(
  //     "mavros/local_position/pose", 1, &cb_pos);
  // ros::Subscriber subLocalVel = nh.subscribe<geometry_msgs::TwistStamped>(
  //     "mavros/local_position/velocity_local", 1, &cb_vel);
  // ros::Subscriber subIMU = nh.subscribe<sensor_msgs::Imu>(
  //     "/mavros/imu/data", 1, &cb_imu);
  ros::Subscriber subOdom = nh.subscribe<nav_msgs::Odometry>(
      "mavros/vision_pose/odom", 1, &cb_odom);
  ros::Subscriber subPVA = nh.subscribe<offboard::ControlCommand>(
      "setpoint_pva", 1, &cb_target_pva);

  /// publisher
  _pubAttThrust = nh.advertise<mavros_msgs::AttitudeTarget>(
      "mavros/setpoint_raw/attitude", 1);
  _pubPID = nh.advertise<offboard::DebugPID>(
      "/Debug/pid", 1);

  ROS_INFO("Subscriber: %s", subState.getTopic().c_str());
  // ROS_INFO("Subscriber: %s", subLocalPos.getTopic().c_str());
  // ROS_INFO("Subscriber: %s", subLocalVel.getTopic().c_str());
  ROS_INFO("Subscriber: %s", subOdom.getTopic().c_str());
  ROS_INFO("Subscriber: %s", subPVA.getTopic().c_str());
  ROS_INFO("Publisher: %s", _pubAttThrust.getTopic().c_str());
  ROS_INFO("Publisher: %s", _pubPID.getTopic().c_str());

  if (!setPIDParam())
  {
    ROS_ERROR("setPIDParam failed , exit !!!");
    return 0;
  }

  ros::Duration(2).sleep();

  // Wait for connection with fcu
  while (ros::ok() && _cur_state.connected == 0)
  {
    ros::spinOnce();
    ros::Duration(1).sleep();
    ROS_INFO("Not Connected");
  }
  ROS_INFO("Connected !!");

  // 临时变量，用于传入当前状态与目标点
  Eigen::Vector3d pos_cur;
  Eigen::Vector3d vel_cur;
  Eigen::Vector3d pos_tar;
  Eigen::Vector3d vel_tar;
  Eigen::Vector3d acc_tar;

  ros::Time t_start = ros::Time::now();
  double t_relative;
  double yaw_tar;
  double yaw_cur;

  while (ros::ok() && _cur_state.connected)
  {
    ros::spinOnce();

    ROS_INFO_THROTTLE(1, "_cur_state: %s", _cur_state.mode.c_str());

    if (!_haveTarget) {
      rate.sleep();
      continue;
    }

    t_relative = (_odom.header.stamp - t_start).toSec();
    pos_cur << _odom.pose.pose.position.x,
        _odom.pose.pose.position.y, _odom.pose.pose.position.z;
    vel_cur << _odom.twist.twist.linear.x,
        _odom.twist.twist.linear.y, _odom.twist.twist.linear.z;

    pos_tar << _pva_msg.pos.x, _pva_msg.pos.y, _pva_msg.pos.z;
    vel_tar << _pva_msg.vel.x, _pva_msg.vel.y, _pva_msg.vel.z;
    acc_tar << _pva_msg.acc.x, _pva_msg.acc.y, _pva_msg.acc.z;

    yaw_tar = _pva_msg.yaw;
    yaw_cur = getYawFromQuaternion(_odom.pose.pose.orientation);

    double yaw_diff = atan2(sin(yaw_tar - yaw_cur), cos(yaw_tar - yaw_cur));
    if (yaw_diff > YAWTHRE)
    {
      if (_cur_state.mode == "OFFBOARD") {
        std::cout << RED_B << "The difference between expected yaw ("
                  << yaw_tar << ") and last yaw (" << yaw_cur
                  << ") is too large" << RESET << endl;
      }

      yaw_tar = yaw_diff > 0 ? yaw_cur + YAWTHRE : yaw_cur - YAWTHRE;
    }

    // 传入PX4串级PID程序解算，并发布控制话题
    px4AttitudeCtlPVA(t_relative,
                      pos_cur, vel_cur,
                      pos_tar, vel_tar, acc_tar, yaw_tar,
                      _cur_state);

    rate.sleep();
  }

  return 0;
}
